#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 使用前需运行：sudo chmod 777 /dev/ttyACM0 以确保串口权限

import rclpy
from rclpy.node import Node
import serial
import time
from geometry_msgs.msg import TwistStamped
from std_msgs.msg import Header

# 定义全局的十六进制数据，用于与力传感器通信
hex_data1 = bytes([0xFE, 0x04, 0x00, 0x00, 0x00, 0x01, 0x25, 0xC5])  # X轴数据请求
hex_data2 = bytes([0xFE, 0x04, 0x00, 0x01, 0x00, 0x01, 0x74, 0x05])  # Y轴数据请求
hex_data3 = bytes([0xFE, 0x04, 0x00, 0x02, 0x00, 0x01, 0x84, 0x05])  # Z轴数据请求
hex_datas = (hex_data1, hex_data2, hex_data3)

class ForceSensorReader(Node):
    def __init__(self):
        # 初始化节点，命名为"force_sensor_reader"
        super().__init__('force_sensor_reader')

        # 初始化串口
        self.ser = serial.Serial(
            port='/dev/ttyACM0',  # 串口设备文件
            baudrate=19200,        # 波特率
            timeout=0.001          # 读取超时时间
        )
        self.ser.flush()  # 清空串口缓冲区

        # 数据滤波和校验相关变量
        self.window_size = 2  # 移动平均滤波窗口大小
        self.x_values = []
        self.y_values = []
        self.z_values = []
        
        # 自动归零功能相关参数
        self.zero_threshold = 0.1  # 归零阈值(N)，小于此值认为是没有外力
        self.stable_count_threshold = 20  # 连续稳定多少次后执行归零
        self.x_stable_count = 0
        self.y_stable_count = 0
        self.z_stable_count = 0
        
        # 上一次的有效值
        self.last_valid_x = 0.0
        self.last_valid_y = 0.0
        self.last_valid_z = 0.0
        
        # 前一次值，用于异常变化检测
        self.prev_x = 0.0
        self.prev_y = 0.0
        self.prev_z = 0.0

        # 创建发布者，发布TwistStamped消息到"/force_data"话题
        self.publisher_ = self.create_publisher(TwistStamped, '/force_data', 10)

        # 创建定时器，每0.1秒调用一次read_force_data函数
        self.timer = self.create_timer(0.01, self.read_force_data)

    def update_filter(self, values, new_value, window_size):
        """更新移动平均滤波器的值"""
        if len(values) >= window_size:
            values.pop(0)
        values.append(new_value)
        return sum(values) / len(values) if values else new_value

    def check_stable(self, current, prev, stable_count, threshold):
        """检查数值是否稳定"""
        if abs(current - prev) < threshold:
            return stable_count + 1
        return 0

    def read_force_data(self):
        # 创建TwistStamped消息对象
        twist_force = TwistStamped()
        twist_force.header = Header()
        twist_force.header.stamp = self.get_clock().now().to_msg()
        twist_force.header.frame_id = 'force_sensor'

        try:
            # 依次读取X、Y、Z轴数据
            for i in range(3):
                retry_count = 0
                data = bytes()
                
                # 增加重试机制
                while retry_count < 3:
                    self.ser.write(hex_datas[i])
                    time.sleep(0.01)
                    data = self.ser.readline()
                    
                    if data and len(data) >= 5:
                        break
                    retry_count += 1

                # 检查数据是否有效
                if data and len(data) >= 5:
                    raw_value = (data[3] << 8 | data[4])
                    
                    # 处理负值
                    if raw_value > 50000:
                        raw_value -= 65535
                    
                    value = raw_value / 100.0  # 转换为浮点数

                    if i == 0:  # X轴数据处理
                        # 异常值检测
                        if abs(value - self.prev_x) > 5.0:
                            # self.get_logger().warn(f"Abnormal X-axis change: {self.prev_x} -> {value}")
                            value = self.prev_x
                        
                        # 更新滤波
                        filtered_x = self.update_filter(self.x_values, value, self.window_size)
                        
                        # 检查稳定性
                        self.x_stable_count = self.check_stable(filtered_x, self.prev_x, self.x_stable_count, 0.05)
                        
                        # 如果稳定且小于阈值，执行归零
                        if self.x_stable_count >= self.stable_count_threshold and abs(filtered_x) < self.zero_threshold:
                            filtered_x = 0.0
                            self.x_values = [0.0] * self.window_size  # 重置滤波器
                        
                        twist_force.twist.linear.x = filtered_x
                        self.prev_x = filtered_x
                        self.last_valid_x = filtered_x

                    elif i == 1:  # Y轴数据处理
                        if abs(value - self.prev_y) > 5.0:
                            self.get_logger().warn(f"Abnormal Y-axis change: {self.prev_y} -> {value}")
                            value = self.prev_y
                            
                        filtered_y = self.update_filter(self.y_values, value, self.window_size)
                        
                        self.y_stable_count = self.check_stable(filtered_y, self.prev_y, self.y_stable_count, 0.05)
                        
                        if self.y_stable_count >= self.stable_count_threshold and abs(filtered_y) < self.zero_threshold:
                            filtered_y = 0.0
                            self.y_values = [0.0] * self.window_size
                            
                        twist_force.twist.linear.y = filtered_y
                        self.prev_y = filtered_y
                        self.last_valid_y = filtered_y

                    elif i == 2:  # Z轴数据处理
                        if abs(value - self.prev_z) > 5.0:
                            self.get_logger().warn(f"Abnormal Z-axis change: {self.prev_z} -> {value}")
                            value = self.prev_z
                            
                        filtered_z = self.update_filter(self.z_values, value, self.window_size)
                        
                        self.z_stable_count = self.check_stable(filtered_z, self.prev_z, self.z_stable_count, 0.05)
                        
                        if self.z_stable_count >= self.stable_count_threshold and abs(filtered_z) < self.zero_threshold:
                            filtered_z = 0.0
                            self.z_values = [0.0] * self.window_size
                            
                        twist_force.twist.linear.z = filtered_z
                        self.prev_z = filtered_z
                        self.last_valid_z = filtered_z
                else:
                    self.get_logger().warn(f"Failed to read data for axis {i} after 3 retries")
                    # 使用上一次的有效值
                    if i == 0:
                        twist_force.twist.linear.x = self.last_valid_x
                    elif i == 1:
                        twist_force.twist.linear.y = self.last_valid_y
                    elif i == 2:
                        twist_force.twist.linear.z = self.last_valid_z

            # 发布消息
            self.publisher_.publish(twist_force)

            # 打印接收到的消息
            self.get_logger().info(
                f"Force data: X={twist_force.twist.linear.x:.2f}, "
                f"Y={twist_force.twist.linear.y:.2f}, "
                f"Z={twist_force.twist.linear.z:.2f} "
                f"(Stable counts: X={self.x_stable_count}, Y={self.y_stable_count}, Z={self.z_stable_count})"
            )

        except Exception as e:
            self.get_logger().error(f"Error reading force data: {e}")
            # 发生异常时发布上一次的有效值
            twist_force.twist.linear.x = self.last_valid_x
            twist_force.twist.linear.y = self.last_valid_y
            twist_force.twist.linear.z = self.last_valid_z
            self.publisher_.publish(twist_force)

def main(args=None):
    rclpy.init(args=args)
    force_sensor_reader = ForceSensorReader()
    rclpy.spin(force_sensor_reader)
    force_sensor_reader.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()